{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8e231c33-44b8-4da2-8d17-90aea2f49ae6",
   "metadata": {},
   "source": [
    "<table style=\"width:100%\">\n",
    "<tr>\n",
    "<td style=\"vertical-align:middle; text-align:left;\">\n",
    "<font size=\"2\">\n",
    "Supplementary code for the <a href=\"https://mng.bz/lZ5B\">Build a Reasoning Model (From Scratch)</a> book by <a href=\"https://sebastianraschka.com\">Sebastian Raschka</a><br>\n",
    "<br>Code repository: <a href=\"https://github.com/rasbt/reasoning-from-scratch\">https://github.com/rasbt/reasoning-from-scratch</a>\n",
    "</font>\n",
    "</td>\n",
    "<td style=\"vertical-align:middle; text-align:left;\">\n",
    "<a href=\"https://mng.bz/lZ5B\"><img src=\"https://sebastianraschka.com/images/reasoning-from-scratch-images/cover-small.webp\" width=\"100px\"></a>\n",
    "</td>\n",
    "</tr>\n",
    "</table>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7bc506a9-fa2f-475e-9981-0c4f36109808",
   "metadata": {},
   "source": [
    "# Chapter 2: Generating Text with a Pre-trained LLM"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78ae290f-bceb-4791-8c38-ef8f7756d1e6",
   "metadata": {},
   "source": [
    "Packages that are being used in this notebook:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f6986129-4aac-496e-a35a-e1dc4c4df4ed",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "reasoning_from_scratch version: 0.1.2\n",
      "torch version: 2.7.1\n",
      "tokenizers version: 0.21.4\n"
     ]
    }
   ],
   "source": [
    "from importlib.metadata import version\n",
    "\n",
    "used_libraries = [\n",
    "    \"reasoning_from_scratch\",\n",
    "    \"torch\",\n",
    "    \"tokenizers\"  # Used by reasoning_from_scratch\n",
    "]\n",
    "\n",
    "for lib in used_libraries:\n",
    "    print(f\"{lib} version: {version(lib)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7fcbc9a-603f-4193-85aa-38dd2d8776bf",
   "metadata": {},
   "source": [
    "<img src=\"https://sebastianraschka.com/images/reasoning-from-scratch-images/ch02/CH02_F01_raschka.webp?1\" width=\"500px\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "378654ef-4323-42f9-8fd4-00f51636e0e1",
   "metadata": {},
   "source": [
    "&nbsp;\n",
    "## 2.1 Introduction to LLMs for text generation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "caba81dd-5a98-4e29-b00d-43b84bd71899",
   "metadata": {},
   "source": [
    "- No code in this section\n",
    "- How do LLMs generate text?\n",
    "- This chapter is a setup chapter: setting up the coding environment and LLM we will be using throughout the book\n",
    "- We also code text generation functions that we will use and extend in upcoming chapters"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8dbef742-7e61-403a-af0a-b95b890c239a",
   "metadata": {},
   "source": [
    "<img src=\"https://sebastianraschka.com/images/reasoning-from-scratch-images/ch02/CH02_F02_raschka.webp?1\" width=\"300px\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0362617f-7976-4507-9308-c2969eab3bba",
   "metadata": {},
   "source": [
    "- LLM (and neural network) flowcharts are traditionally read and drawn from top to bottom"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "186d56ba-6c74-4d69-a14f-d5bbae7852fe",
   "metadata": {},
   "source": [
    "&nbsp;\n",
    "## 2.2 Setting up the coding environment"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dcc92122-0ce4-4426-8e85-4a48e1ebe593",
   "metadata": {},
   "source": [
    "- If you are reading this book, you likely coded in Python before\n",
    "- The simplest way to install dependencies, if you already have a Python environment set up (with Python 3.10 or newer), is to use `pip`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "eaf19556-05a5-47c3-8031-cb2f4a2fac2c",
   "metadata": {},
   "outputs": [],
   "source": [
    "#!pip install -r https://raw.githubusercontent.com/rasbt/reasoning-from-scratch/refs/heads/main/requirements.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0f86ccc-395c-415e-a631-71e9e6ec59a4",
   "metadata": {},
   "source": [
    "- For this chapter, dependencies can also be installed manually:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "cabb6da5-891f-4aac-b478-c2e8ea28b0d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "#!pip install torch>=2.7.1 tokenizers>=0.21.2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0236c753-d983-49e4-b8cb-3870ecdae943",
   "metadata": {},
   "source": [
    "- My preferred way is to use the widely recommended [uv](https://docs.astral.sh/uv/) Python package and project manager\n",
    "- To install `uv`, run the installation for your OS from the official website: https://docs.astral.sh/uv/getting-started/installation/\n",
    "- Next, clone the GitHub repo:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "efa5b432-536c-45f4-b73d-568d9788b934",
   "metadata": {},
   "outputs": [],
   "source": [
    "#!git clone --depth 1 https://github.com/rasbt/reasoning-from-scratch.git"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "670db5b1-66be-4382-ae8c-2f7e0189ac2e",
   "metadata": {},
   "source": [
    "- If you don't have `git` installed, you can also manually download the source code repository from the Manning website or by clicking this link: https://github.com/rasbt/reasoning-from-scratch/archive/refs/heads/main.zip (unzip it after downloading)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f3ea394-bf43-4a0a-ad04-df3a82b9874b",
   "metadata": {},
   "source": [
    "- In the terminal, navigate to the `reasoning-from-scratch` folder\n",
    "- Run `uv run jupyter lab` to launch JupyterLab and open a blank notebook or the notebook for this chapter\n",
    "- This command also sets up a local virtual environment (usually in `.venv/`) and installs all dependencies from the `pyproject.toml` file inside the `reasoning-from-scratch` folder automatically"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4373b49-ebf3-4282-9f21-7ac32ed5607b",
   "metadata": {},
   "source": [
    "<img src=\"https://sebastianraschka.com/images/reasoning-from-scratch-images/ch02/CH02_F03_raschka.webp?1\" width=\"500px\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0c28e8a-d328-4ff2-9a55-38af8b56ffae",
   "metadata": {},
   "source": [
    "- See [../02_setup-tips/python-instructions.md](../02_setup-tips/python-instructions.md) for additional installation details and options if needed"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b16265a-7cd2-406b-b954-b5aaebf1e09b",
   "metadata": {},
   "source": [
    "&nbsp;\n",
    "## 2.3 Understanding hardware needs and recommendations"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "331d182d-1747-48ae-a502-004bf8730bf9",
   "metadata": {},
   "source": [
    "- If you are new to PyTorch, I recommend reading through my [PyTorch in One Hour: From Tensors to Training Neural Networks on Multiple GPUs](https://sebastianraschka.com/teaching/pytorch-1h/) tutorial\n",
    "- If you followed the previous section, you should have PyTorch installed\n",
    "- Check manually if your PyTorch installation supports GPU; see what's supported on your machine: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "8d205144-19bf-481f-be5b-c1e77847c025",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PyTorch version 2.7.1\n",
      "Apple Silicon GPU\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "print(f\"PyTorch version {torch.__version__}\")\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    print(\"CUDA GPU\")\n",
    "elif torch.mps.is_available():\n",
    "    print(\"Apple Silicon GPU\")\n",
    "else:\n",
    "    print(\"Only CPU\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a08b8fc-3fc6-4440-9cd5-be4c9ed7d43b",
   "metadata": {},
   "source": [
    "- Depending on the chapter, code will automatically use NVIDIA GPU if available, otherwise run on CPU (or Apple Silicon GPU if recommended for a particular section or chapter)\n",
    "- Chapters 2-4 can be executed in a reasonable time on a CPU\n",
    "- Code in chapters 5-7 will be very slow when executed on a CPU, and a GPU with NVIDIA is recommended for these chapters (more on the exact resource needs in those upcoming chapters)\n",
    "- My personal preference is [Lightning AI Studio](https://lightning.ai/), which offers users free compute credits after the sign-up and verification process; alternatively, [Google Colab](https://colab.research.google.com/) is another good choice"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7fd7329f-5b6b-4001-afcc-ddde1778689f",
   "metadata": {},
   "source": [
    "<img src=\"https://sebastianraschka.com/images/reasoning-from-scratch-images/ch02/CH02_F04_raschka.webp\" width=\"500px\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a7b5c47-d22b-4cce-9f82-64be1ca2ce57",
   "metadata": {},
   "source": [
    "- See [../02_setup-tips/gpu-instructions.md](../02_setup-tips/gpu-instructions.md) for cloud compute recommendations if needed\n",
    "- But for now, there is no need to use GPUs yet; the first chapters run fine on non-GPU hardware"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f846ae6-c6ca-44a5-8c19-c5b79728e8c3",
   "metadata": {},
   "source": [
    "&nbsp;\n",
    "## 2.4 Preparing input texts for LLMs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2dfd4adc-9774-4813-873d-a456ff129f5c",
   "metadata": {},
   "source": [
    "- In this section, we learn how to use a tokenizer; we use it to convert (encode) input text into a token ID representation as input to the LLM\n",
    "- We also use the tokenizer to convert (decode) the LLM output back into a human-readable text representation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56135ed3-a423-43b3-ab00-75fc3d19e2d6",
   "metadata": {},
   "source": [
    "<img src=\"https://sebastianraschka.com/images/reasoning-from-scratch-images/ch02/CH02_F05_raschka.webp?1\" width=\"500px\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2123a05e-91b3-42ad-a468-8d6694a49c45",
   "metadata": {},
   "source": [
    "- As mentioned earlier, implementing the LLM and tokenizer from scratch is outside the scope of this book, which is focused on implementing reasoning methods from scratch on top of an existing LLM and tokenizer\n",
    "- In this book, we will work with a pre-trained LLM that we will load in the next section; here, we load the tokenizer that goes with it\n",
    "- I prepared a `reasoning_from_scratch` Python package that provides the base LLM and the corresponding tokenizer, which I coded with the help of the [`tokenizers`](https://github.com/huggingface/tokenizers) Python library package\n",
    "- The `reasoning_from_scratch` package code is part of this book's supplementary code, and it should already be installed based on the instructions in section 2.2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31f5ae78-5aed-477d-8c1e-4fc1df857834",
   "metadata": {},
   "source": [
    "- Next, we download the tokenizer files (this is a tokenizer for the Qwen3 base LLM, but more on that in the next section):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "6a9aba9e-0fa8-4710-804c-a26299beb2f5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "✓ qwen3/tokenizer-base.json already up-to-date\n"
     ]
    }
   ],
   "source": [
    "from reasoning_from_scratch.qwen3 import download_qwen3_small\n",
    "\n",
    "download_qwen3_small(kind=\"base\", tokenizer_only=True, out_dir=\"qwen3\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "079987a0-b0dd-4dc2-8432-10be34029990",
   "metadata": {},
   "source": [
    "- Now, we can load the tokenizer settings from the tokenizer file into the `Qwen3Tokenizer`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "564206d6-8180-44de-813f-3bf7290d59f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "from reasoning_from_scratch.qwen3 import Qwen3Tokenizer\n",
    "\n",
    "tokenizer_file_path = Path(\"qwen3\") / \"tokenizer-base.json\"\n",
    "tokenizer = Qwen3Tokenizer(tokenizer_file_path=tokenizer_file_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4e7efec-1fda-430b-b06e-4f856ec3be13",
   "metadata": {},
   "source": [
    "- Since we haven't loaded the LLM itself yet, we will do a simpler round-trip: we encode the text into token IDs and then encode it back into its string representation:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2435704e-765a-4cb0-8960-25314cb96e4d",
   "metadata": {},
   "source": [
    "<img src=\"https://sebastianraschka.com/images/reasoning-from-scratch-images/ch02/CH02_F06_raschka.webp\" width=\"500px\">"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "6c6c4fb1-2c54-45e2-b30d-22717527fffc",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = \"Explain large language models.\"\n",
    "input_token_ids_list = tokenizer.encode(prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "c914b8cd-5905-4fae-92f6-df9aecf3811b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "840 --> Ex\n",
      "20772 --> plain\n",
      "3460 -->  large\n",
      "4128 -->  language\n",
      "4119 -->  models\n",
      "13 --> .\n"
     ]
    }
   ],
   "source": [
    "for i in input_token_ids_list:\n",
    "    print(f\"{i} --> {tokenizer.decode([i])}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "94d58966-cee7-4ed4-ba41-80d064dd4fd8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Explain large language models.\n"
     ]
    }
   ],
   "source": [
    "text = tokenizer.decode(input_token_ids_list)\n",
    "print(text)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33683b20-c8ad-434d-8eeb-7fa779378238",
   "metadata": {},
   "source": [
    "- In case of the `Qwen3Tokenizer`, there are about 151 thousand unique tokens (vocabulary size)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b48b9bb-99d7-4688-ad37-79596ca89701",
   "metadata": {},
   "source": [
    "- Additional resources on tokenization:\n",
    "  - [Build a Large Language Model (from Scratch)](https://mng.bz/M96o) chapter 2\n",
    "  - [Implementing A Byte Pair Encoding (BPE) Tokenizer From Scratch](https://sebastianraschka.com/blog/2025/bpe-from-scratch.html)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93295da8-8c52-49fe-93c8-5daa8e8c679c",
   "metadata": {},
   "source": [
    "&nbsp;\n",
    "## 2.5 Loading pre-trained models"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe267761-51f1-49eb-b36e-f415f367d831",
   "metadata": {},
   "source": [
    "<img src=\"https://sebastianraschka.com/images/reasoning-from-scratch-images/ch02/CH02_F07_raschka.webp\" width=\"500px\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a31c2558-45d7-4d62-af94-1d91f869f8a2",
   "metadata": {},
   "source": [
    "- As hinted at in the previous section, when loading the tokenizer, this book uses Qwen3 0.6B; after thinking long and hard about which open weight base model to use, I opted for Qwen3 because\n",
    "  - Qwen3 is the leading open-weight model in terms of modeling performance as of this writing\n",
    "  - Qwen3 0.6B is more memory efficient than Llama 3 1B\n",
    "  - There's both a base model (which we focus on for reasoning model development) and an official reasoning variant that we can use as a reference model\n",
    "- (Note that the canonical spelling does not include a whitespace in \"Qwen3\" whereas it includes one in \"Llama 3\")\n",
    "- In the spirit of \"from-scratch\" we are using a reimplementation of Qwen3 that I wrote in pure PyTorch without any external LLM library dependencies; this from-scratch implementation is compatible with the original Qwen3 model weights\n",
    "- However, we will not go over the Qwen3 code implementation in this book as this would be a whole book by itself (similar to my [Build A Large Language Model (From Scratch)](https://github.com/rasbt/LLMs-from-scratch) book; instead, this book (Build A Reasoning Model From Scratch) focuses on implementing reasoning methods from scratch on top of a base model (here Qwen3)\n",
    "- See appendix C for the Qwen3 model code\n",
    "- See appendix D for loading the reasoning variant and larger Qwen3 models\n",
    "- See the Qwen3 [GitHub repository](https://github.com/QwenLM/Qwen3) and [technical report](https://arxiv.org/abs/2505.09388) for (even) more details"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3afd498b-38f4-4430-9a52-0eec8d288c88",
   "metadata": {},
   "source": [
    "- The model is purposefully small (but still very capable) to run on consumer hardware\n",
    "- It runs fine on CPU, NVIDIA GPUs (`\"cuda\"`), and Apple Silicon GPUs (`\"mps\"`); more about the performance trade-offs later in this chapter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "afbc8f36-b76f-481c-b291-83abd6fb20ad",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using Apple Silicon GPU (MPS)\n"
     ]
    }
   ],
   "source": [
    "def get_device():\n",
    "    if torch.cuda.is_available():\n",
    "        device = torch.device(\"cuda\")\n",
    "        print(\"Using NVIDIA CUDA GPU\")\n",
    "    elif torch.backends.mps.is_available():\n",
    "        device = torch.device(\"mps\")\n",
    "        print(\"Using Apple Silicon GPU (MPS)\")\n",
    "    else:\n",
    "        device = torch.device(\"cpu\")\n",
    "        print(\"Using CPU\")\n",
    "    return device\n",
    "\n",
    "device = get_device()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6068984f-bff3-4189-b843-25002573a2e6",
   "metadata": {},
   "source": [
    "- I recommend running the code on `\"cpu\"` on the first run-through, so we hardcode the device below: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "542875a3-cd13-4b66-a5a8-510817fe4d3b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Recommended: Use CPU on the first run-through\n",
    "device = torch.device(\"cpu\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a76c6f6a-f215-4248-ac5f-8bc4212e288e",
   "metadata": {},
   "source": [
    "- Then, we download the file containing the pre-trained model weights, which is approximately 1.5 GB in size:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "52145acb-24b6-44fb-963b-33d9fe9398e3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "✓ qwen3/qwen3-0.6B-base.pth already up-to-date\n",
      "✓ qwen3/tokenizer-base.json already up-to-date\n"
     ]
    }
   ],
   "source": [
    "download_qwen3_small(kind=\"base\", tokenizer_only=False, out_dir=\"qwen3\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef90039d-2a46-48c7-9d7f-51822d1ad18f",
   "metadata": {},
   "source": [
    "- The architectural structure of the Qwen3 0.6B model we are loading is shown below for readers who are familiar with LLM architectures, but note that for this book, it's **not** essential or important to understand this architecture as we are not modifying but rather adding reasoning techniques on top in later chapters"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f02be0d-e6b0-4cb3-8533-8111dd9313b5",
   "metadata": {},
   "source": [
    "- I coded the Qwen3 model architecture from scratch for the [reasoning-from-scratch](https://github.com/rasbt/reasoning-from-scratch/blob/main/reasoning_from_scratch/qwen3.py) Python package contained in this code repository; the source code is also shown in appendix C; but again, this is only as a bonus for those who are curious, and it's not necessary to look at or understand these internals to follow the rest of the book"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "501d6703-5fbf-4254-8b86-010bb0244b75",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Qwen3Model(\n",
       "  (tok_emb): Embedding(151936, 1024)\n",
       "  (trf_blocks): ModuleList(\n",
       "    (0-27): 28 x TransformerBlock(\n",
       "      (att): GroupedQueryAttention(\n",
       "        (W_query): Linear(in_features=1024, out_features=2048, bias=False)\n",
       "        (W_key): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "        (W_value): Linear(in_features=1024, out_features=1024, bias=False)\n",
       "        (out_proj): Linear(in_features=2048, out_features=1024, bias=False)\n",
       "        (q_norm): RMSNorm()\n",
       "        (k_norm): RMSNorm()\n",
       "      )\n",
       "      (ff): FeedForward(\n",
       "        (fc1): Linear(in_features=1024, out_features=3072, bias=False)\n",
       "        (fc2): Linear(in_features=1024, out_features=3072, bias=False)\n",
       "        (fc3): Linear(in_features=3072, out_features=1024, bias=False)\n",
       "      )\n",
       "      (norm1): RMSNorm()\n",
       "      (norm2): RMSNorm()\n",
       "    )\n",
       "  )\n",
       "  (final_norm): RMSNorm()\n",
       "  (out_head): Linear(in_features=1024, out_features=151936, bias=False)\n",
       ")"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from reasoning_from_scratch.qwen3 import Qwen3Model, QWEN_CONFIG_06_B\n",
    "\n",
    "model_file = Path(\"qwen3\") / \"qwen3-0.6B-base.pth\"\n",
    "\n",
    "model = Qwen3Model(QWEN_CONFIG_06_B)\n",
    "model.load_state_dict(torch.load(model_file))\n",
    "\n",
    "model.to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da0583d6-74ef-4629-b436-6f64f8c0841b",
   "metadata": {},
   "source": [
    "<img src=\"https://sebastianraschka.com/images/reasoning-from-scratch-images/ch02/CH02_F08_raschka.webp\" width=\"300px\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d2c7723-2e57-4d1a-8a10-3b4daafe4d50",
   "metadata": {},
   "source": [
    "&nbsp;\n",
    "## 2.6 Understanding the sequential LLM text generation process"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b38845f2-55e9-4e4a-a8a5-d402d7450501",
   "metadata": {},
   "source": [
    "- In this section, we code a simple wrapper function so we can use the LLM to generate text (we will extend this function with extra functionality in chapter 4)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4dd89a18-de70-4958-8d24-d0aaddc519f3",
   "metadata": {},
   "source": [
    "<img src=\"https://sebastianraschka.com/images/reasoning-from-scratch-images/ch02/CH02_F09_raschka.webp?1\" width=\"500px\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "446fdb75-6971-4bad-bfde-4f8dab1e5315",
   "metadata": {},
   "source": [
    "- LLMs generate one word at a time:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9cee173-54bc-4275-8fcc-320c1a1a38fb",
   "metadata": {},
   "source": [
    "<img src=\"https://sebastianraschka.com/images/reasoning-from-scratch-images/ch02/CH02_F10_raschka.webp?2\" width=\"500px\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61bb43dc-9007-48b3-bd7b-432c7b616aba",
   "metadata": {},
   "source": [
    "- The figure above is a simplification, only showing the newly generated word; the figure below zooms in on the first iteration:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e78cff76-1215-4d10-9c3c-3a5f2afbb987",
   "metadata": {},
   "source": [
    "<img src=\"https://sebastianraschka.com/images/reasoning-from-scratch-images/ch02/CH02_F11_raschka.webp\" width=\"3b00px\">"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "94ed4960-8e72-48ae-90ab-662b468f77a2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([1, 2, 3])\n",
      "tensor([[1, 2, 3]])\n"
     ]
    }
   ],
   "source": [
    "example = torch.tensor([1, 2, 3]) \n",
    "print(example)\n",
    "print(example.unsqueeze(0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "1d65f427-6cb2-4533-85e6-0b09f9226e0c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[1, 2, 3]])\n",
      "tensor([1, 2, 3])\n"
     ]
    }
   ],
   "source": [
    "example = torch.tensor([[1, 2, 3]]) \n",
    "print(example)\n",
    "print(example.squeeze(0))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4cbc1181-eca7-4a25-8586-d6a9f08fa3f0",
   "metadata": {},
   "source": [
    "<img src=\"https://sebastianraschka.com/images/reasoning-from-scratch-images/ch02/CH02_F11_raschka.webp?2\" width=\"300px\">"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "f2c44846-dd0a-4877-bc23-7baa96ff50c8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of input tokens: 6\n",
      "Formatted Output tensor shape: torch.Size([6, 151936])\n"
     ]
    }
   ],
   "source": [
    "prompt = \"Explain large language models.\"\n",
    "input_token_ids_list = tokenizer.encode(prompt)\n",
    "print(f\"Number of input tokens: {len(input_token_ids_list)}\")\n",
    "\n",
    "input_tensor = torch.tensor(input_token_ids_list)\n",
    "input_tensor_fmt = input_tensor.unsqueeze(0)\n",
    "\n",
    "output_tensor = model(input_tensor_fmt)\n",
    "output_tensor_fmt = output_tensor.squeeze(0)\n",
    "print(f\"Formatted Output tensor shape: {output_tensor_fmt.shape}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7de206fa-5f1c-4356-8c77-3c0365a7f1c6",
   "metadata": {},
   "source": [
    "<img src=\"https://sebastianraschka.com/images/reasoning-from-scratch-images/ch02/CH02_F12_raschka.webp\" width=\"500px\">"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "72c10d27-f29e-4fa6-a748-7e3096b25979",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([ 7.3750,  2.0312,  8.0000,  ..., -2.5469, -2.5469, -2.5469],\n",
      "       dtype=torch.bfloat16)\n"
     ]
    }
   ],
   "source": [
    "last_token = output_tensor_fmt[-1].detach()\n",
    "print(last_token)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "7fa3b6ab-daee-46ed-b476-20feaa284a4f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([20286])\n"
     ]
    }
   ],
   "source": [
    "print(last_token.argmax(dim=-1, keepdim=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "a705ea38-2bc7-430a-8fe7-e8b58e994dd4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Large\n"
     ]
    }
   ],
   "source": [
    "print(tokenizer.decode([20286]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "008ad321-9e3b-428f-9dba-89aa86738870",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(3)\n",
      "tensor(2)\n"
     ]
    }
   ],
   "source": [
    "example = torch.tensor([-2, 1, 3, 1])\n",
    "print(torch.max(example))\n",
    "print(torch.argmax(example))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "05c87e6d-d77c-4f4a-b265-0d0a95d684c1",
   "metadata": {},
   "source": [
    "&nbsp;\n",
    "## 2.7 Coding a minimal text generation function\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4e89807-92c0-4462-a447-d9c45bbf02ca",
   "metadata": {},
   "source": [
    "<img src=\"https://sebastianraschka.com/images/reasoning-from-scratch-images/ch02/CH02_F13_raschka.webp\" width=\"500px\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ec19437-a3f0-45b7-b572-d7a520165ff1",
   "metadata": {},
   "source": [
    "<img src=\"https://sebastianraschka.com/images/reasoning-from-scratch-images/ch02/CH02_F14_raschka.webp?2\" width=\"500px\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4dea54fe-7292-4901-bc99-62064bba48ee",
   "metadata": {},
   "source": [
    "- The `generate_text_basic` function implements this sequential text generation process:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "a0c1b331-79f6-4fad-a56f-8aaf8b2a1239",
   "metadata": {},
   "outputs": [],
   "source": [
    "@torch.inference_mode()\n",
    "def generate_text_basic(\n",
    "    model,\n",
    "    token_ids,\n",
    "    max_new_tokens,\n",
    "    eos_token_id=None\n",
    "):\n",
    "    input_length = token_ids.shape[1]\n",
    "    model.eval()\n",
    "\n",
    "    for _ in range(max_new_tokens):\n",
    "        out = model(token_ids)[:, -1]\n",
    "        next_token = torch.argmax(out, dim=-1, keepdim=True)\n",
    "\n",
    "        # Stop if all sequences in the batch have generated EOS\n",
    "        if (eos_token_id is not None\n",
    "                and torch.all(next_token == eos_token_id)):\n",
    "            break\n",
    "\n",
    "        token_ids = torch.cat([token_ids, next_token], dim=1)\n",
    "    return token_ids[:, input_length:]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74f14a54-fd55-4bc5-a834-7c2ccd3f30ae",
   "metadata": {},
   "source": [
    "- Let's use it to generate a 100-token response to a simple \"Explain large language models in 2 sentences.\" prompt to see how it works (we get to the reasoning parts in later chapters)\n",
    "- The following code will be slow and can take 1-3 minutes to complete, depending on your computer (we will speed it up in later sections) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "4edc105a-0ba4-4bf4-b4f6-b66de0bdea65",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Large language models are artificial intelligence systems that can understand, generate, and process human language, enabling them to perform a wide range of tasks, from answering questions to writing articles, and even creating creative content.<|endoftext|>Human language is a complex and dynamic system that has evolved over millions of years to enable effective communication and social interaction. It is composed of a vast array of symbols, including letters, numbers, and words, which are used to convey meaning and express thoughts and ideas. The evolution of language has\n"
     ]
    }
   ],
   "source": [
    "prompt = \"Explain large language models in a single sentence.\"\n",
    "input_token_ids_tensor = torch.tensor(\n",
    "    tokenizer.encode(prompt),\n",
    "    device=device\n",
    "    ).unsqueeze(0)\n",
    "\n",
    "max_new_tokens = 100\n",
    "output_token_ids_tensor = generate_text_basic(\n",
    "    model=model,\n",
    "    token_ids=input_token_ids_tensor,\n",
    "    max_new_tokens=max_new_tokens,\n",
    ")\n",
    "output_text = tokenizer.decode(\n",
    "    output_token_ids_tensor.squeeze(0).tolist()\n",
    ")\n",
    "print(output_text)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3489d09f-1178-4870-a7ee-4faabf507717",
   "metadata": {},
   "source": [
    "- Notice that the LLM follows the instruction quite well, but the response becomes nonsensical/off-topic after `<|endoftext|>`, which is a token used as a delimiter between different documents during training\n",
    "- When using the LLM, we want it to stop generating after encountering this token"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "38e676de-a07e-4eb6-90d0-9f8f10da9caf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[151643]\n"
     ]
    }
   ],
   "source": [
    "print(tokenizer.encode(\"<|endoftext|>\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f146f8b-4ed1-4939-9d3a-3de81e8039e0",
   "metadata": {},
   "source": [
    "- For convenience, this token ID is stored as a tokenizer attribute (eos = end of sequence):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "8a547b6e-9154-49be-9b0b-c4fcc3da462d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "151643\n"
     ]
    }
   ],
   "source": [
    "print(tokenizer.eos_token_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "676290fe-1d2e-41db-8ca2-717fe8cb2561",
   "metadata": {},
   "source": [
    "- We can use it to tell the LLM (or rather the `generate_text_basic` function) when to stop generating text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "4c97d10c-e3de-4fcd-a669-3c4198be4018",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Large language models are artificial intelligence systems that can understand, generate, and process human language, enabling them to perform a wide range of tasks, from answering questions to writing articles, and even creating creative content.\n"
     ]
    }
   ],
   "source": [
    "output_token_ids_tensor = generate_text_basic(\n",
    "    model=model,\n",
    "    token_ids=input_token_ids_tensor,\n",
    "    max_new_tokens=max_new_tokens,\n",
    "    eos_token_id=tokenizer.eos_token_id\n",
    ")\n",
    "\n",
    "output_text = tokenizer.decode(\n",
    "    output_token_ids_tensor.squeeze(0).tolist()\n",
    ")\n",
    "print(output_text)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0150d6f6-e6cc-4576-b77a-7f12b53d3aa6",
   "metadata": {},
   "source": [
    "- The response above is what you get when running to code on CPU, the generated text may differ slightly differ depending on the device"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5669f8e-b6e9-490f-90cb-52735b3aab1a",
   "metadata": {},
   "source": [
    "- Before we wrap up this section and see how we can speed up the code, let's implement a simple benchmarking function to track the computational performance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "9efe08e3-6819-490f-bbfa-88b4e61b4468",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_stats(output_token_ids, tokenizer, start_time, end_time):\n",
    "    total_time = end_time - start_time\n",
    "    print(f\"Time: {total_time:.2f} sec\")\n",
    "    print(f\"{int(output_token_ids.numel() / total_time)} tokens/sec\")\n",
    "\n",
    "    if torch.cuda.is_available():\n",
    "        max_mem_bytes = torch.cuda.max_memory_allocated()\n",
    "        max_mem_gb = max_mem_bytes / (1024 ** 3)\n",
    "        print(f\"Max memory allocated: {max_mem_gb:.2f} GB\")\n",
    "\n",
    "    output_text = tokenizer.decode(output_token_ids.squeeze(0).tolist())\n",
    "    print(f\"\\n{output_text}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "1e66c303-d4e9-4a9b-9f6f-ecfe90e6e3c1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time: 9.23 sec\n",
      "4 tokens/sec\n",
      "\n",
      " Large language models are artificial intelligence systems that can understand, generate, and process human language, enabling them to perform a wide range of tasks, from answering questions to writing articles, and even creating creative content.\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "\n",
    "start_time = time.time()\n",
    "output_token_ids_tensor = generate_text_basic(\n",
    "    model=model,\n",
    "    token_ids=input_token_ids_tensor,\n",
    "    max_new_tokens=max_new_tokens,\n",
    "    eos_token_id=tokenizer.eos_token_id\n",
    ")\n",
    "end_time = time.time()\n",
    "\n",
    "\n",
    "generate_stats(output_token_ids_tensor, tokenizer, start_time, end_time)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce1c6163-e78d-4f03-8ec3-b2d8af767564",
   "metadata": {},
   "source": [
    "&nbsp;\n",
    "## 2.8 Faster inference via KV caching"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ba8fe63-1dc8-4992-b4d1-ab994160abb3",
   "metadata": {},
   "source": [
    "<img src=\"https://sebastianraschka.com/images/reasoning-from-scratch-images/ch02/CH02_F15_raschka.webp?2\" width=\"500px\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56dbd7d4-f6ed-43ab-9751-8534e06dfcc7",
   "metadata": {},
   "source": [
    "- Note that the code in this book emphasizes code readability, and a whole separate book can be written about optimizations\n",
    "- Here, we look at an engineering trick called \"KV caching\" (KV refers to the keys and values inside the attention mechanism of the LLM)\n",
    "- If you are unfamiliar with these terms, don't worry, all you need to know is that there is a way we can store (cache) intermediate values that are reused in each iteration"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2253ebb3-3125-4cde-b562-a0d8bf662888",
   "metadata": {},
   "source": [
    "<img src=\"https://sebastianraschka.com/images/reasoning-from-scratch-images/ch02/CH02_F16_raschka.webp\" width=\"500px\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02471c83-8237-48f6-adff-2811152ca203",
   "metadata": {},
   "source": [
    "- For more details on the mechanics of KV caching, see my [Understanding and Coding the KV Cache in LLMs from Scratch](https://magazine.sebastianraschka.com/p/coding-the-kv-cache-in-llms) article\n",
    "- Below is a modified version of the `generate_text_basic` function that uses a KV cache"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "d4095088-9dd5-4ffd-b2f5-24d848ae27f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from reasoning_from_scratch.qwen3 import KVCache\n",
    "\n",
    "@torch.inference_mode()\n",
    "def generate_text_basic_cache(\n",
    "    model,\n",
    "    token_ids,\n",
    "    max_new_tokens,\n",
    "    eos_token_id=None\n",
    "):\n",
    "\n",
    "    input_length = token_ids.shape[1] \n",
    "    model.eval()\n",
    "    cache = KVCache(n_layers=model.cfg[\"n_layers\"])\n",
    "    model.reset_kv_cache()\n",
    "\n",
    "    out = model(token_ids, cache=cache)[:, -1]\n",
    "    for _ in range(max_new_tokens):\n",
    "        next_token = torch.argmax(out, dim=-1, keepdim=True)\n",
    "\n",
    "        if (eos_token_id is not None \n",
    "               and torch.all(next_token == eos_token_id)):\n",
    "            break\n",
    "\n",
    "        token_ids = torch.cat([token_ids, next_token], dim=1)\n",
    "        out = model(next_token, cache=cache)[:, -1]\n",
    "\n",
    "    return token_ids[:, input_length:]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad5e6c37-fafe-4143-a88d-959916e4f301",
   "metadata": {},
   "source": [
    "- The usage is similar to before:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "d8027190-6de5-4d6f-ae3a-2baf2027b2f0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time: 1.40 sec\n",
      "29 tokens/sec\n",
      "\n",
      " Large language models are artificial intelligence systems that can understand, generate, and process human language, enabling them to perform a wide range of tasks, from answering questions to writing articles, and even creating creative content.\n"
     ]
    }
   ],
   "source": [
    "start_time = time.time()\n",
    "\n",
    "output_token_ids_tensor = generate_text_basic_cache(\n",
    "    model=model,\n",
    "    token_ids=input_token_ids_tensor,\n",
    "    max_new_tokens=max_new_tokens,\n",
    "    eos_token_id=tokenizer.eos_token_id,\n",
    ")\n",
    "end_time = time.time()\n",
    "\n",
    "generate_stats(output_token_ids_tensor, tokenizer, start_time, end_time)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4235f50d-c8cc-4e53-bcf1-abd8dffac617",
   "metadata": {},
   "source": [
    "- As we can see, it is magnitudes faster than before (28 tokens/sec instead of 4 tokens/sec; run on a Mac Mini M4 CPU)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f4e58ad-df61-4f5f-8582-9378e8ba7ed0",
   "metadata": {},
   "source": [
    "&nbsp;\n",
    "## 2.9 Faster inference via PyTorch model compilation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb5e1ded-e4a9-47a3-bf68-7f7241d878e5",
   "metadata": {},
   "source": [
    "<img src=\"https://sebastianraschka.com/images/reasoning-from-scratch-images/ch02/CH02_F17_raschka.webp?2\" width=\"500px\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12863930-7bb9-4d0d-8a79-5d5582a5a38e",
   "metadata": {},
   "source": [
    "- Another technique to speed up the model inference (text generation) by a lot is using `torch.compile`\n",
    "- Note that this currently doesn't work on MPS (Apple Silicon GPU) devices due to `InductorError`\n",
    "- The usage is simple, we just call `torch.compile` on the model (see [the documentation](https://docs.pytorch.org/docs/stable/torch.compiler.html) for additional options)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "a23ac9c1-4e83-4b05-b5b2-093178450931",
   "metadata": {},
   "outputs": [],
   "source": [
    "if device.type == \"mps\":\n",
    "    print(f\"`torch.compile` is not supported for the {model.__class__.__name__} model on MPS (Apple Silicon) as of this writing.\")\n",
    "    model_compiled = model\n",
    "else:\n",
    "    model_compiled = torch.compile(model)\n",
    "    # Assignment so that notebook doesn't stop here if someone uses \"Run All Cells\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d179612a-2f46-44d5-b180-385eb3bfadc9",
   "metadata": {},
   "source": [
    "- The first iteration can be a bit slow as it does the initial compilation and optimization; hence, we repeat the text generation multiple times\n",
    "- First, let's start with the non-cached version (this can be a bit slow and might take xx minutes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "43044619-36a7-4036-a7a6-c0172ea7630e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Warm-up run\n",
      "Time: 27.70 sec\n",
      "1 tokens/sec\n",
      "\n",
      " Large language models are artificial intelligence systems that can understand, generate, and process human language, enabling them to perform a wide range of tasks, from answering questions to writing articles, and even creating creative content.\n",
      "\n",
      "------------------------------\n",
      "\n",
      "Timed run 1:\n",
      "Time: 7.09 sec\n",
      "5 tokens/sec\n",
      "\n",
      " Large language models are artificial intelligence systems that can understand, generate, and process human language, enabling them to perform a wide range of tasks, from answering questions to writing articles, and even creating creative content.\n",
      "\n",
      "------------------------------\n",
      "\n",
      "Timed run 2:\n",
      "Time: 7.19 sec\n",
      "5 tokens/sec\n",
      "\n",
      " Large language models are artificial intelligence systems that can understand, generate, and process human language, enabling them to perform a wide range of tasks, from answering questions to writing articles, and even creating creative content.\n",
      "\n",
      "------------------------------\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for i in range(3):\n",
    "    start_time = time.time()\n",
    "    output_token_ids_tensor = generate_text_basic(\n",
    "        model=model_compiled,\n",
    "        token_ids=input_token_ids_tensor,\n",
    "        max_new_tokens=max_new_tokens,\n",
    "        eos_token_id=tokenizer.eos_token_id\n",
    "    )\n",
    "    end_time = time.time()\n",
    "\n",
    "    if i == 0:\n",
    "        print(\"Warm-up run\")\n",
    "    else:\n",
    "        print(f\"Timed run {i}:\")\n",
    "    generate_stats(output_token_ids_tensor, tokenizer, start_time, end_time)\n",
    "\n",
    "    print(f\"\\n{30*'-'}\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d86b44a9-0512-4c74-b699-aa7e4b68cc0c",
   "metadata": {},
   "source": [
    "- As we can see above, with 5 tokens/sec, this is only marginally faster than before (4 tokens/sec)\n",
    "- Let's now see how well the KV cache version does"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "397bd803-a554-46a1-a2fd-eb8d5e65b134",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Warm-up run\n",
      "Time: 29.87 sec\n",
      "1 tokens/sec\n",
      "\n",
      " Large language models are artificial intelligence systems that can understand, generate, and process human language, enabling them to perform a wide range of tasks, from answering questions to writing articles, and even creating creative content.\n",
      "\n",
      "------------------------------\n",
      "\n",
      "Timed run 1:\n",
      "Time: 0.60 sec\n",
      "68 tokens/sec\n",
      "\n",
      " Large language models are artificial intelligence systems that can understand, generate, and process human language, enabling them to perform a wide range of tasks, from answering questions to writing articles, and even creating creative content.\n",
      "\n",
      "------------------------------\n",
      "\n",
      "Timed run 2:\n",
      "Time: 0.62 sec\n",
      "66 tokens/sec\n",
      "\n",
      " Large language models are artificial intelligence systems that can understand, generate, and process human language, enabling them to perform a wide range of tasks, from answering questions to writing articles, and even creating creative content.\n",
      "\n",
      "------------------------------\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for i in range(3):\n",
    "    start_time = time.time()\n",
    "    output_token_ids_tensor = generate_text_basic_cache(\n",
    "        model=model_compiled,\n",
    "        token_ids=input_token_ids_tensor,\n",
    "        max_new_tokens=max_new_tokens,\n",
    "        eos_token_id=tokenizer.eos_token_id\n",
    "    )\n",
    "    end_time = time.time()\n",
    "\n",
    "    if i == 0:\n",
    "        print(\"Warm-up run\")\n",
    "    else:\n",
    "        print(f\"Timed run {i}:\")\n",
    "    generate_stats(\n",
    "        output_token_ids_tensor, tokenizer, start_time, end_time\n",
    "    )\n",
    "\n",
    "    print(f\"\\n{30*'-'}\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4c98ca9-a9fd-4f82-957f-eaae1e6907a7",
   "metadata": {},
   "source": [
    "- As we can see, the compilation resulted in a substantial 2x speed-up (64 tokens/sec versus 30 tokens/sec)\n",
    "- Below is a table with additional results"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e403cdac-c633-49a6-a713-7735efc46a60",
   "metadata": {},
   "source": [
    "| Model      | Mode              | Hardware        | Tokens/sec    | GPU Memory (VRAM) |\n",
    "|------------|-------------------|-----------------|---------------|-------------------|\n",
    "| Qwen3Model | Regular           | Mac Mini M4 CPU | 6             | -                 |\n",
    "| Qwen3Model | Regular compiled  | Mac Mini M4 CPU | 6             | -                 |\n",
    "| Qwen3Model | KV cache          | Mac Mini M4 CPU | 28            | -                 |\n",
    "| Qwen3Model | KV cache compiled | Mac Mini M4 CPU | 68            | -                 |\n",
    "|            |                   |                 |               |                   |\n",
    "| Qwen3Model | Regular           | Mac Mini M4 GPU | 17            | -                 |\n",
    "| Qwen3Model | Regular compiled  | Mac Mini M4 GPU | InductorError | -                 |\n",
    "| Qwen3Model | KV cache          | Mac Mini M4 GPU | 18            | -                 |\n",
    "| Qwen3Model | KV cache compiled | Mac Mini M4 GPU | InductorError | -                 |\n",
    "|            |                   |                 |               |                   |\n",
    "| Qwen3Model | Regular           | NVIDIA H100 GPU | 51            | 1.55 GB           |\n",
    "| Qwen3Model | Regular compiled  | NVIDIA H100 GPU | 164           | 1.81 GB           |\n",
    "| Qwen3Model | KV cache          | NVIDIA H100 GPU | 48            | 1.52 GB           |\n",
    "| Qwen3Model | KV cache compiled | NVIDIA H100 GPU | 141           | 1.81 GB           |"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d05c6a6-3170-4251-9bdf-07489d1104c6",
   "metadata": {},
   "source": [
    "- Note that we ran all the examples with a single prompt (i.e., a batch size of 1); if you are curious about batched inference, see appendix E"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0db3a7b3-9fb2-47de-9f42-4ad1f1f25de8",
   "metadata": {},
   "source": [
    "&nbsp;\n",
    "## Summary"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40a8d136-06df-46d4-b616-202dd9a8cdeb",
   "metadata": {},
   "source": [
    "- No code in this section"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
